{
 "cells": [
  {
   "cell_type": "raw",
   "id": "b61d713f-3ef8-4bef-8ff1-9859276c7726",
   "metadata": {},
   "source": [
    "[2024-09-25]\n",
    "PTB-XL 데이터 전처리 진행\n",
    "process_ecg.py 파일 그대로 jupyter notebook에서 실행하는 방향으로 !!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0a969593-842a-469d-8ce8-c61518679c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import wfdb\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "# _LEAD_NAMES = [\"I\", \"II\", \"III\", \"aVR\", \"aVL\", \"aVF\", \"V1\", \"V2\", \"V3\", \"V4\", \"V5\", \"V6\"]\n",
    "_LEAD_NAMES = ['I', 'II', 'III', 'AVR', 'AVL', 'AVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']"
   ]
  },
  {
   "cell_type": "raw",
   "id": "fb898884-9912-4e53-b6f7-564962cbe1d8",
   "metadata": {},
   "source": [
    "if __name__ == \"__main__\":\n",
    "    run(get_parser())\n",
    "\n",
    "python process_ecg.py \\\n",
    "    --input_dir ${DATABASE_DIRECTORY} \\\n",
    "    --output_dir ${WAVEFORM_DIRECTORY} \\\n",
    "    --index_path ${INDEX_PATH}\n",
    "\n",
    "- input_dir: 원시 WFDB 파형 파일이 저장된 directory\n",
    "- output_dir: ECG PKL 파일 저장될 directory\n",
    "- index_path: index csv 파일 저장될 경로"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "319dfb1f-5ac8-46d3-8d3a-a9038939a83b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_parser():\n",
    "    description = \"Process WFDB ECG database.\"\n",
    "    # parser = argparse.ArgumentParser(description=description)\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('-i',\n",
    "                        '--input_dir',\n",
    "                        type=str,\n",
    "                        # required=True,\n",
    "                        default='/tf/physionet.org/files/ptb-xl/1.0.3/records500',\n",
    "                        help=\"Path to the WFDB ECG database directory.\")\n",
    "    parser.add_argument('-o',\n",
    "                        '--output_dir',\n",
    "                        type=str,\n",
    "                        # required=True,\n",
    "                        default='./ptbxl/ecgs500/',\n",
    "                        help=\"Path to the directory where the preprocessed signals will be saved.\")\n",
    "    parser.add_argument('--index_path',\n",
    "                        type=str,\n",
    "                        default='./ptbxl/index.csv',\n",
    "                        help=\"Path to the index file.\")\n",
    "    args = parser.parse_args(\"\")\n",
    "    return args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "41ae9fa0-ad53-4fa9-a5c6-e9163bddd7a2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Namespace(input_dir='/tf/physionet.org/files/ptb-xl/1.0.3/records500', output_dir='./ptbxl/ecgs500/', index_path='./ptbxl/index.csv')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args = get_parser()\n",
    "args"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53e06139-70d8-4e27-a8d7-3ccf531599ee",
   "metadata": {},
   "source": [
    "get_parser() 부분 실행 완 => args 처리됨. 이후 run 코드 순서대로 실행하면 됨!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ea547cef-aab4-487c-8911-56393c721ad4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_records(root_dir):\n",
    "    \"\"\"Find all the .hea files in the root directory and its subdirectories.\n",
    "    Args:\n",
    "        root_dir (str): The directory to search for .hea files.\n",
    "    Returns:\n",
    "        records (set): A set of record names.\n",
    "                       (e.g., ['database/1/ecg001', 'database/1/ecg001', ..., 'database/9/ecg991'])\n",
    "    \"\"\"\n",
    "    records = set()\n",
    "    for root, _, files in os.walk(root_dir):\n",
    "        for file in files:\n",
    "            extension = os.path.splitext(file)[1]\n",
    "            if extension == '.hea':\n",
    "                record = os.path.relpath(os.path.join(root, file), root_dir)[:-4]\n",
    "                records.add(record)\n",
    "    records = sorted(records)\n",
    "    return records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4a348247-59b8-4609-b0b9-c19bf09fd3f3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['00000/00001_hr',\n",
       " '00000/00002_hr',\n",
       " '00000/00003_hr',\n",
       " '00000/00004_hr',\n",
       " '00000/00005_hr',\n",
       " '00000/00006_hr',\n",
       " '00000/00007_hr',\n",
       " '00000/00008_hr',\n",
       " '00000/00009_hr',\n",
       " '00000/00010_hr',\n",
       " '00000/00011_hr',\n",
       " '00000/00012_hr',\n",
       " '00000/00013_hr',\n",
       " '00000/00014_hr',\n",
       " '00000/00015_hr',\n",
       " '00000/00016_hr',\n",
       " '00000/00017_hr',\n",
       " '00000/00018_hr',\n",
       " '00000/00019_hr',\n",
       " '00000/00020_hr',\n",
       " '00000/00021_hr',\n",
       " '00000/00022_hr',\n",
       " '00000/00023_hr',\n",
       " '00000/00024_hr',\n",
       " '00000/00025_hr',\n",
       " '00000/00026_hr',\n",
       " '00000/00027_hr',\n",
       " '00000/00028_hr',\n",
       " '00000/00029_hr',\n",
       " '00000/00030_hr',\n",
       " '00000/00031_hr',\n",
       " '00000/00032_hr',\n",
       " '00000/00033_hr',\n",
       " '00000/00034_hr',\n",
       " '00000/00035_hr',\n",
       " '00000/00036_hr',\n",
       " '00000/00037_hr',\n",
       " '00000/00038_hr',\n",
       " '00000/00039_hr',\n",
       " '00000/00040_hr',\n",
       " '00000/00041_hr',\n",
       " '00000/00042_hr',\n",
       " '00000/00043_hr',\n",
       " '00000/00044_hr',\n",
       " '00000/00045_hr',\n",
       " '00000/00046_hr',\n",
       " '00000/00047_hr',\n",
       " '00000/00048_hr',\n",
       " '00000/00049_hr',\n",
       " '00000/00050_hr',\n",
       " '00000/00051_hr',\n",
       " '00000/00052_hr',\n",
       " '00000/00053_hr',\n",
       " '00000/00054_hr',\n",
       " '00000/00055_hr',\n",
       " '00000/00056_hr',\n",
       " '00000/00057_hr',\n",
       " '00000/00058_hr',\n",
       " '00000/00059_hr',\n",
       " '00000/00060_hr',\n",
       " '00000/00061_hr',\n",
       " '00000/00062_hr',\n",
       " '00000/00063_hr',\n",
       " '00000/00064_hr',\n",
       " '00000/00065_hr',\n",
       " '00000/00066_hr',\n",
       " '00000/00067_hr',\n",
       " '00000/00068_hr',\n",
       " '00000/00069_hr',\n",
       " '00000/00070_hr',\n",
       " '00000/00071_hr',\n",
       " '00000/00072_hr',\n",
       " '00000/00073_hr',\n",
       " '00000/00074_hr',\n",
       " '00000/00075_hr',\n",
       " '00000/00076_hr',\n",
       " '00000/00077_hr',\n",
       " '00000/00078_hr',\n",
       " '00000/00079_hr',\n",
       " '00000/00080_hr',\n",
       " '00000/00081_hr',\n",
       " '00000/00082_hr',\n",
       " '00000/00083_hr',\n",
       " '00000/00084_hr',\n",
       " '00000/00085_hr',\n",
       " '00000/00086_hr',\n",
       " '00000/00087_hr',\n",
       " '00000/00088_hr',\n",
       " '00000/00089_hr',\n",
       " '00000/00090_hr',\n",
       " '00000/00091_hr',\n",
       " '00000/00092_hr',\n",
       " '00000/00093_hr',\n",
       " '00000/00094_hr',\n",
       " '00000/00095_hr',\n",
       " '00000/00096_hr',\n",
       " '00000/00097_hr',\n",
       " '00000/00098_hr',\n",
       " '00000/00099_hr',\n",
       " '00000/00100_hr',\n",
       " '00000/00101_hr',\n",
       " '00000/00102_hr',\n",
       " '00000/00103_hr',\n",
       " '00000/00104_hr',\n",
       " '00000/00105_hr',\n",
       " '00000/00106_hr',\n",
       " '00000/00107_hr',\n",
       " '00000/00108_hr',\n",
       " '00000/00109_hr',\n",
       " '00000/00110_hr',\n",
       " '00000/00111_hr',\n",
       " '00000/00112_hr',\n",
       " '00000/00113_hr',\n",
       " '00000/00114_hr',\n",
       " '00000/00115_hr',\n",
       " '00000/00116_hr',\n",
       " '00000/00117_hr',\n",
       " '00000/00118_hr',\n",
       " '00000/00119_hr',\n",
       " '00000/00120_hr',\n",
       " '00000/00121_hr',\n",
       " '00000/00122_hr',\n",
       " '00000/00123_hr',\n",
       " '00000/00124_hr',\n",
       " '00000/00125_hr',\n",
       " '00000/00126_hr',\n",
       " '00000/00127_hr',\n",
       " '00000/00128_hr',\n",
       " '00000/00129_hr',\n",
       " '00000/00130_hr',\n",
       " '00000/00131_hr',\n",
       " '00000/00132_hr',\n",
       " '00000/00133_hr',\n",
       " '00000/00134_hr',\n",
       " '00000/00135_hr',\n",
       " '00000/00136_hr',\n",
       " '00000/00138_hr',\n",
       " '00000/00144_hr',\n",
       " '00000/00146_hr',\n",
       " '00000/00147_hr',\n",
       " '00000/00148_hr',\n",
       " '00000/00149_hr',\n",
       " '00000/00150_hr',\n",
       " '00000/00151_hr',\n",
       " '00000/00152_hr',\n",
       " '00000/00153_hr',\n",
       " '00000/00154_hr',\n",
       " '00000/00155_hr',\n",
       " '00000/00156_hr',\n",
       " '00000/00157_hr',\n",
       " '00000/00158_hr',\n",
       " '00000/00159_hr',\n",
       " '00000/00160_hr',\n",
       " '00000/00161_hr',\n",
       " '00000/00162_hr',\n",
       " '00000/00163_hr',\n",
       " '00000/00164_hr',\n",
       " '00000/00165_hr',\n",
       " '00000/00166_hr',\n",
       " '00000/00167_hr',\n",
       " '00000/00168_hr',\n",
       " '00000/00169_hr',\n",
       " '00000/00170_hr',\n",
       " '00000/00171_hr',\n",
       " '00000/00172_hr',\n",
       " '00000/00173_hr',\n",
       " '00000/00174_hr',\n",
       " '00000/00175_hr',\n",
       " '00000/00176_hr',\n",
       " '00000/00177_hr',\n",
       " '00000/00178_hr',\n",
       " '00000/00179_hr',\n",
       " '00000/00180_hr',\n",
       " '00000/00181_hr',\n",
       " '00000/00182_hr',\n",
       " '00000/00183_hr',\n",
       " '00000/00184_hr',\n",
       " '00000/00185_hr',\n",
       " '00000/00186_hr',\n",
       " '00000/00187_hr',\n",
       " '00000/00188_hr',\n",
       " '00000/00189_hr',\n",
       " '00000/00190_hr',\n",
       " '00000/00191_hr',\n",
       " '00000/00192_hr',\n",
       " '00000/00193_hr',\n",
       " '00000/00194_hr',\n",
       " '00000/00195_hr',\n",
       " '00000/00196_hr',\n",
       " '00000/00197_hr',\n",
       " '00000/00198_hr',\n",
       " '00000/00199_hr',\n",
       " '00000/00200_hr',\n",
       " '00000/00201_hr',\n",
       " '00000/00202_hr',\n",
       " '00000/00203_hr',\n",
       " '00000/00204_hr',\n",
       " '00000/00205_hr',\n",
       " '00000/00206_hr',\n",
       " '00000/00207_hr',\n",
       " '00000/00208_hr',\n",
       " '00000/00209_hr',\n",
       " '00000/00210_hr',\n",
       " '00000/00211_hr',\n",
       " '00000/00212_hr',\n",
       " '00000/00213_hr',\n",
       " '00000/00214_hr',\n",
       " '00000/00215_hr',\n",
       " '00000/00216_hr',\n",
       " '00000/00217_hr',\n",
       " '00000/00218_hr',\n",
       " '00000/00219_hr',\n",
       " '00000/00220_hr',\n",
       " '00000/00221_hr',\n",
       " '00000/00222_hr',\n",
       " '00000/00223_hr',\n",
       " '00000/00224_hr',\n",
       " '00000/00225_hr',\n",
       " '00000/00226_hr',\n",
       " '00000/00227_hr',\n",
       " '00000/00228_hr',\n",
       " '00000/00229_hr',\n",
       " '00000/00230_hr',\n",
       " '00000/00231_hr',\n",
       " '00000/00232_hr',\n",
       " '00000/00233_hr',\n",
       " '00000/00234_hr',\n",
       " '00000/00235_hr',\n",
       " '00000/00236_hr',\n",
       " '00000/00237_hr',\n",
       " '00000/00238_hr',\n",
       " '00000/00239_hr',\n",
       " '00000/00240_hr',\n",
       " '00000/00241_hr',\n",
       " '00000/00242_hr',\n",
       " '00000/00243_hr',\n",
       " '00000/00244_hr',\n",
       " '00000/00245_hr',\n",
       " '00000/00246_hr',\n",
       " '00000/00247_hr',\n",
       " '00000/00248_hr',\n",
       " '00000/00249_hr',\n",
       " '00000/00250_hr',\n",
       " '00000/00251_hr',\n",
       " '00000/00252_hr',\n",
       " '00000/00253_hr',\n",
       " '00000/00254_hr',\n",
       " '00000/00255_hr',\n",
       " '00000/00256_hr',\n",
       " '00000/00257_hr',\n",
       " '00000/00258_hr',\n",
       " '00000/00259_hr',\n",
       " '00000/00260_hr',\n",
       " '00000/00261_hr',\n",
       " '00000/00262_hr',\n",
       " '00000/00263_hr',\n",
       " '00000/00264_hr',\n",
       " '00000/00265_hr',\n",
       " '00000/00266_hr',\n",
       " '00000/00267_hr',\n",
       " '00000/00268_hr',\n",
       " '00000/00269_hr',\n",
       " '00000/00270_hr',\n",
       " '00000/00271_hr',\n",
       " '00000/00272_hr',\n",
       " '00000/00273_hr',\n",
       " '00000/00274_hr',\n",
       " '00000/00275_hr',\n",
       " '00000/00276_hr',\n",
       " '00000/00277_hr',\n",
       " '00000/00278_hr',\n",
       " '00000/00279_hr',\n",
       " '00000/00280_hr',\n",
       " '00000/00281_hr',\n",
       " '00000/00282_hr',\n",
       " '00000/00283_hr',\n",
       " '00000/00284_hr',\n",
       " '00000/00285_hr',\n",
       " '00000/00286_hr',\n",
       " '00000/00287_hr',\n",
       " '00000/00288_hr',\n",
       " '00000/00289_hr',\n",
       " '00000/00290_hr',\n",
       " '00000/00291_hr',\n",
       " '00000/00292_hr',\n",
       " '00000/00293_hr',\n",
       " '00000/00294_hr',\n",
       " '00000/00295_hr',\n",
       " '00000/00296_hr',\n",
       " '00000/00297_hr',\n",
       " '00000/00298_hr',\n",
       " '00000/00299_hr',\n",
       " '00000/00300_hr',\n",
       " '00000/00301_hr',\n",
       " '00000/00302_hr',\n",
       " '00000/00303_hr',\n",
       " '00000/00304_hr',\n",
       " '00000/00305_hr',\n",
       " '00000/00306_hr',\n",
       " '00000/00307_hr',\n",
       " '00000/00308_hr',\n",
       " '00000/00309_hr',\n",
       " '00000/00310_hr',\n",
       " '00000/00311_hr',\n",
       " '00000/00312_hr',\n",
       " '00000/00313_hr',\n",
       " '00000/00314_hr',\n",
       " '00000/00315_hr',\n",
       " '00000/00316_hr',\n",
       " '00000/00317_hr',\n",
       " '00000/00318_hr',\n",
       " '00000/00319_hr',\n",
       " '00000/00320_hr',\n",
       " '00000/00321_hr',\n",
       " '00000/00322_hr',\n",
       " '00000/00323_hr',\n",
       " '00000/00324_hr',\n",
       " '00000/00325_hr',\n",
       " '00000/00326_hr',\n",
       " '00000/00327_hr',\n",
       " '00000/00328_hr',\n",
       " '00000/00329_hr',\n",
       " '00000/00330_hr',\n",
       " '00000/00331_hr',\n",
       " '00000/00332_hr',\n",
       " '00000/00333_hr',\n",
       " '00000/00334_hr',\n",
       " '00000/00335_hr',\n",
       " '00000/00336_hr',\n",
       " '00000/00337_hr',\n",
       " '00000/00338_hr',\n",
       " '00000/00339_hr',\n",
       " '00000/00340_hr',\n",
       " '00000/00341_hr',\n",
       " '00000/00342_hr',\n",
       " '00000/00343_hr',\n",
       " '00000/00344_hr',\n",
       " '00000/00345_hr',\n",
       " '00000/00346_hr',\n",
       " '00000/00347_hr',\n",
       " '00000/00348_hr',\n",
       " '00000/00349_hr',\n",
       " '00000/00350_hr',\n",
       " '00000/00351_hr',\n",
       " '00000/00352_hr',\n",
       " '00000/00353_hr',\n",
       " '00000/00354_hr',\n",
       " '00000/00355_hr',\n",
       " '00000/00356_hr',\n",
       " '00000/00357_hr',\n",
       " '00000/00358_hr',\n",
       " '00000/00359_hr',\n",
       " '00000/00360_hr',\n",
       " '00000/00361_hr',\n",
       " '00000/00362_hr',\n",
       " '00000/00363_hr',\n",
       " '00000/00364_hr',\n",
       " '00000/00365_hr',\n",
       " '00000/00366_hr',\n",
       " '00000/00367_hr',\n",
       " '00000/00368_hr',\n",
       " '00000/00369_hr',\n",
       " '00000/00370_hr',\n",
       " '00000/00371_hr',\n",
       " '00000/00372_hr',\n",
       " '00000/00373_hr',\n",
       " '00000/00374_hr',\n",
       " '00000/00375_hr',\n",
       " '00000/00376_hr',\n",
       " '00000/00377_hr',\n",
       " '00000/00378_hr',\n",
       " '00000/00379_hr',\n",
       " '00000/00380_hr',\n",
       " '00000/00381_hr',\n",
       " '00000/00382_hr',\n",
       " '00000/00383_hr',\n",
       " '00000/00384_hr',\n",
       " '00000/00385_hr',\n",
       " '00000/00386_hr',\n",
       " '00000/00387_hr',\n",
       " '00000/00388_hr',\n",
       " '00000/00389_hr',\n",
       " '00000/00390_hr',\n",
       " '00000/00391_hr',\n",
       " '00000/00392_hr',\n",
       " '00000/00393_hr',\n",
       " '00000/00394_hr',\n",
       " '00000/00395_hr',\n",
       " '00000/00396_hr',\n",
       " '00000/00397_hr',\n",
       " '00000/00398_hr',\n",
       " '00000/00399_hr',\n",
       " '00000/00400_hr',\n",
       " '00000/00401_hr',\n",
       " '00000/00402_hr',\n",
       " '00000/00403_hr',\n",
       " '00000/00404_hr',\n",
       " '00000/00405_hr',\n",
       " '00000/00406_hr',\n",
       " '00000/00407_hr',\n",
       " '00000/00408_hr',\n",
       " '00000/00409_hr',\n",
       " '00000/00410_hr',\n",
       " '00000/00411_hr',\n",
       " '00000/00412_hr',\n",
       " '00000/00413_hr',\n",
       " '00000/00414_hr',\n",
       " '00000/00415_hr',\n",
       " '00000/00416_hr',\n",
       " '00000/00417_hr',\n",
       " '00000/00418_hr',\n",
       " '00000/00419_hr',\n",
       " '00000/00420_hr',\n",
       " '00000/00421_hr',\n",
       " '00000/00422_hr',\n",
       " '00000/00423_hr',\n",
       " '00000/00424_hr',\n",
       " '00000/00425_hr',\n",
       " '00000/00426_hr',\n",
       " '00000/00427_hr',\n",
       " '00000/00428_hr',\n",
       " '00000/00429_hr',\n",
       " '00000/00430_hr',\n",
       " '00000/00431_hr',\n",
       " '00000/00432_hr',\n",
       " '00000/00433_hr',\n",
       " '00000/00434_hr',\n",
       " '00000/00435_hr',\n",
       " '00000/00436_hr',\n",
       " '00000/00437_hr',\n",
       " '00000/00438_hr',\n",
       " '00000/00439_hr',\n",
       " '00000/00440_hr',\n",
       " '00000/00441_hr',\n",
       " '00000/00442_hr',\n",
       " '00000/00443_hr',\n",
       " '00000/00444_hr',\n",
       " '00000/00445_hr',\n",
       " '00000/00446_hr',\n",
       " '00000/00447_hr',\n",
       " '00000/00448_hr',\n",
       " '00000/00449_hr',\n",
       " '00000/00450_hr',\n",
       " '00000/00451_hr',\n",
       " '00000/00452_hr',\n",
       " '00000/00453_hr',\n",
       " '00000/00454_hr',\n",
       " '00000/00455_hr',\n",
       " '00000/00457_hr',\n",
       " '00000/00460_hr',\n",
       " '00000/00463_hr',\n",
       " '00000/00464_hr',\n",
       " '00000/00465_hr',\n",
       " '00000/00466_hr',\n",
       " '00000/00467_hr',\n",
       " '00000/00468_hr',\n",
       " '00000/00469_hr',\n",
       " '00000/00470_hr',\n",
       " '00000/00471_hr',\n",
       " '00000/00472_hr',\n",
       " '00000/00473_hr',\n",
       " '00000/00474_hr',\n",
       " '00000/00475_hr',\n",
       " '00000/00476_hr',\n",
       " '00000/00477_hr',\n",
       " '00000/00478_hr',\n",
       " '00000/00479_hr',\n",
       " '00000/00480_hr',\n",
       " '00000/00481_hr',\n",
       " '00000/00482_hr',\n",
       " '00000/00483_hr',\n",
       " '00000/00484_hr',\n",
       " '00000/00485_hr',\n",
       " '00000/00486_hr',\n",
       " '00000/00487_hr',\n",
       " '00000/00488_hr',\n",
       " '00000/00489_hr',\n",
       " '00000/00490_hr',\n",
       " '00000/00491_hr',\n",
       " '00000/00492_hr',\n",
       " '00000/00493_hr',\n",
       " '00000/00494_hr',\n",
       " '00000/00495_hr',\n",
       " '00000/00496_hr',\n",
       " '00000/00497_hr',\n",
       " '00000/00498_hr',\n",
       " '00000/00499_hr',\n",
       " '00000/00500_hr',\n",
       " '00000/00501_hr',\n",
       " '00000/00502_hr',\n",
       " '00000/00503_hr',\n",
       " '00000/00504_hr',\n",
       " '00000/00505_hr',\n",
       " '00000/00506_hr',\n",
       " '00000/00507_hr',\n",
       " '00000/00508_hr',\n",
       " '00000/00509_hr',\n",
       " '00000/00510_hr',\n",
       " '00000/00511_hr',\n",
       " '00000/00512_hr',\n",
       " '00000/00513_hr',\n",
       " '00000/00514_hr',\n",
       " '00000/00515_hr',\n",
       " '00000/00516_hr',\n",
       " '00000/00517_hr',\n",
       " '00000/00518_hr',\n",
       " '00000/00519_hr',\n",
       " '00000/00520_hr',\n",
       " '00000/00521_hr',\n",
       " '00000/00522_hr',\n",
       " '00000/00523_hr',\n",
       " '00000/00524_hr',\n",
       " '00000/00525_hr',\n",
       " '00000/00526_hr',\n",
       " '00000/00527_hr',\n",
       " '00000/00528_hr',\n",
       " '00000/00529_hr',\n",
       " '00000/00530_hr',\n",
       " '00000/00531_hr',\n",
       " '00000/00532_hr',\n",
       " '00000/00533_hr',\n",
       " '00000/00534_hr',\n",
       " '00000/00535_hr',\n",
       " '00000/00536_hr',\n",
       " '00000/00537_hr',\n",
       " '00000/00538_hr',\n",
       " '00000/00539_hr',\n",
       " '00000/00540_hr',\n",
       " '00000/00541_hr',\n",
       " '00000/00542_hr',\n",
       " '00000/00543_hr',\n",
       " '00000/00544_hr',\n",
       " '00000/00545_hr',\n",
       " '00000/00546_hr',\n",
       " '00000/00547_hr',\n",
       " '00000/00548_hr',\n",
       " '00000/00549_hr',\n",
       " '00000/00550_hr',\n",
       " '00000/00551_hr',\n",
       " '00000/00552_hr',\n",
       " '00000/00553_hr',\n",
       " '00000/00554_hr',\n",
       " '00000/00555_hr',\n",
       " '00000/00556_hr',\n",
       " '00000/00557_hr',\n",
       " '00000/00558_hr',\n",
       " '00000/00559_hr',\n",
       " '00000/00560_hr',\n",
       " '00000/00561_hr',\n",
       " '00000/00562_hr',\n",
       " '00000/00563_hr',\n",
       " '00000/00564_hr',\n",
       " '00000/00565_hr',\n",
       " '00000/00566_hr',\n",
       " '00000/00567_hr',\n",
       " '00000/00568_hr',\n",
       " '00000/00569_hr',\n",
       " '00000/00570_hr',\n",
       " '00000/00571_hr',\n",
       " '00000/00572_hr',\n",
       " '00000/00573_hr',\n",
       " '00000/00574_hr',\n",
       " '00000/00575_hr',\n",
       " '00000/00576_hr',\n",
       " '00000/00577_hr',\n",
       " '00000/00578_hr',\n",
       " '00000/00579_hr',\n",
       " '00000/00580_hr',\n",
       " '00000/00581_hr',\n",
       " '00000/00582_hr',\n",
       " '00000/00583_hr',\n",
       " '00000/00584_hr',\n",
       " '00000/00585_hr',\n",
       " '00000/00586_hr',\n",
       " '00000/00587_hr',\n",
       " '00000/00588_hr',\n",
       " '00000/00589_hr',\n",
       " '00000/00590_hr',\n",
       " '00000/00591_hr',\n",
       " '00000/00592_hr',\n",
       " '00000/00593_hr',\n",
       " '00000/00594_hr',\n",
       " '00000/00595_hr',\n",
       " '00000/00596_hr',\n",
       " '00000/00597_hr',\n",
       " '00000/00598_hr',\n",
       " '00000/00599_hr',\n",
       " '00000/00600_hr',\n",
       " '00000/00601_hr',\n",
       " '00000/00602_hr',\n",
       " '00000/00603_hr',\n",
       " '00000/00604_hr',\n",
       " '00000/00605_hr',\n",
       " '00000/00606_hr',\n",
       " '00000/00607_hr',\n",
       " '00000/00608_hr',\n",
       " '00000/00609_hr',\n",
       " '00000/00610_hr',\n",
       " '00000/00611_hr',\n",
       " '00000/00612_hr',\n",
       " '00000/00613_hr',\n",
       " '00000/00614_hr',\n",
       " '00000/00615_hr',\n",
       " '00000/00616_hr',\n",
       " '00000/00617_hr',\n",
       " '00000/00618_hr',\n",
       " '00000/00619_hr',\n",
       " '00000/00620_hr',\n",
       " '00000/00621_hr',\n",
       " '00000/00622_hr',\n",
       " '00000/00623_hr',\n",
       " '00000/00624_hr',\n",
       " '00000/00625_hr',\n",
       " '00000/00626_hr',\n",
       " '00000/00627_hr',\n",
       " '00000/00628_hr',\n",
       " '00000/00629_hr',\n",
       " '00000/00630_hr',\n",
       " '00000/00631_hr',\n",
       " '00000/00632_hr',\n",
       " '00000/00633_hr',\n",
       " '00000/00634_hr',\n",
       " '00000/00635_hr',\n",
       " '00000/00636_hr',\n",
       " '00000/00637_hr',\n",
       " '00000/00638_hr',\n",
       " '00000/00639_hr',\n",
       " '00000/00640_hr',\n",
       " '00000/00641_hr',\n",
       " '00000/00642_hr',\n",
       " '00000/00643_hr',\n",
       " '00000/00644_hr',\n",
       " '00000/00645_hr',\n",
       " '00000/00646_hr',\n",
       " '00000/00647_hr',\n",
       " '00000/00648_hr',\n",
       " '00000/00649_hr',\n",
       " '00000/00650_hr',\n",
       " '00000/00651_hr',\n",
       " '00000/00652_hr',\n",
       " '00000/00653_hr',\n",
       " '00000/00654_hr',\n",
       " '00000/00655_hr',\n",
       " '00000/00656_hr',\n",
       " '00000/00657_hr',\n",
       " '00000/00658_hr',\n",
       " '00000/00659_hr',\n",
       " '00000/00660_hr',\n",
       " '00000/00661_hr',\n",
       " '00000/00662_hr',\n",
       " '00000/00663_hr',\n",
       " '00000/00664_hr',\n",
       " '00000/00665_hr',\n",
       " '00000/00666_hr',\n",
       " '00000/00667_hr',\n",
       " '00000/00668_hr',\n",
       " '00000/00669_hr',\n",
       " '00000/00670_hr',\n",
       " '00000/00671_hr',\n",
       " '00000/00672_hr',\n",
       " '00000/00673_hr',\n",
       " '00000/00674_hr',\n",
       " '00000/00675_hr',\n",
       " '00000/00676_hr',\n",
       " '00000/00677_hr',\n",
       " '00000/00678_hr',\n",
       " '00000/00679_hr',\n",
       " '00000/00680_hr',\n",
       " '00000/00681_hr',\n",
       " '00000/00682_hr',\n",
       " '00000/00683_hr',\n",
       " '00000/00684_hr',\n",
       " '00000/00685_hr',\n",
       " '00000/00686_hr',\n",
       " '00000/00687_hr',\n",
       " '00000/00688_hr',\n",
       " '00000/00689_hr',\n",
       " '00000/00690_hr',\n",
       " '00000/00691_hr',\n",
       " '00000/00692_hr',\n",
       " '00000/00693_hr',\n",
       " '00000/00694_hr',\n",
       " '00000/00695_hr',\n",
       " '00000/00696_hr',\n",
       " '00000/00697_hr',\n",
       " '00000/00698_hr',\n",
       " '00000/00699_hr',\n",
       " '00000/00700_hr',\n",
       " '00000/00701_hr',\n",
       " '00000/00702_hr',\n",
       " '00000/00703_hr',\n",
       " '00000/00704_hr',\n",
       " '00000/00705_hr',\n",
       " '00000/00706_hr',\n",
       " '00000/00707_hr',\n",
       " '00000/00708_hr',\n",
       " '00000/00709_hr',\n",
       " '00000/00710_hr',\n",
       " '00000/00711_hr',\n",
       " '00000/00712_hr',\n",
       " '00000/00713_hr',\n",
       " '00000/00714_hr',\n",
       " '00000/00715_hr',\n",
       " '00000/00716_hr',\n",
       " '00000/00717_hr',\n",
       " '00000/00718_hr',\n",
       " '00000/00719_hr',\n",
       " '00000/00720_hr',\n",
       " '00000/00721_hr',\n",
       " '00000/00722_hr',\n",
       " '00000/00723_hr',\n",
       " '00000/00724_hr',\n",
       " '00000/00725_hr',\n",
       " '00000/00726_hr',\n",
       " '00000/00727_hr',\n",
       " '00000/00728_hr',\n",
       " '00000/00729_hr',\n",
       " '00000/00730_hr',\n",
       " '00000/00731_hr',\n",
       " '00000/00732_hr',\n",
       " '00000/00733_hr',\n",
       " '00000/00734_hr',\n",
       " '00000/00735_hr',\n",
       " '00000/00736_hr',\n",
       " '00000/00737_hr',\n",
       " '00000/00738_hr',\n",
       " '00000/00739_hr',\n",
       " '00000/00740_hr',\n",
       " '00000/00741_hr',\n",
       " '00000/00742_hr',\n",
       " '00000/00743_hr',\n",
       " '00000/00744_hr',\n",
       " '00000/00745_hr',\n",
       " '00000/00746_hr',\n",
       " '00000/00747_hr',\n",
       " '00000/00748_hr',\n",
       " '00000/00749_hr',\n",
       " '00000/00750_hr',\n",
       " '00000/00751_hr',\n",
       " '00000/00752_hr',\n",
       " '00000/00753_hr',\n",
       " '00000/00754_hr',\n",
       " '00000/00755_hr',\n",
       " '00000/00756_hr',\n",
       " '00000/00757_hr',\n",
       " '00000/00758_hr',\n",
       " '00000/00759_hr',\n",
       " '00000/00760_hr',\n",
       " '00000/00761_hr',\n",
       " '00000/00762_hr',\n",
       " '00000/00763_hr',\n",
       " '00000/00764_hr',\n",
       " '00000/00765_hr',\n",
       " '00000/00766_hr',\n",
       " '00000/00767_hr',\n",
       " '00000/00768_hr',\n",
       " '00000/00769_hr',\n",
       " '00000/00770_hr',\n",
       " '00000/00771_hr',\n",
       " '00000/00772_hr',\n",
       " '00000/00773_hr',\n",
       " '00000/00774_hr',\n",
       " '00000/00775_hr',\n",
       " '00000/00776_hr',\n",
       " '00000/00777_hr',\n",
       " '00000/00778_hr',\n",
       " '00000/00779_hr',\n",
       " '00000/00780_hr',\n",
       " '00000/00781_hr',\n",
       " '00000/00782_hr',\n",
       " '00000/00783_hr',\n",
       " '00000/00784_hr',\n",
       " '00000/00785_hr',\n",
       " '00000/00786_hr',\n",
       " '00000/00787_hr',\n",
       " '00000/00788_hr',\n",
       " '00000/00789_hr',\n",
       " '00000/00790_hr',\n",
       " '00000/00791_hr',\n",
       " '00000/00792_hr',\n",
       " '00000/00793_hr',\n",
       " '00000/00794_hr',\n",
       " '00000/00795_hr',\n",
       " '00000/00796_hr',\n",
       " '00000/00797_hr',\n",
       " '00000/00798_hr',\n",
       " '00000/00799_hr',\n",
       " '00000/00800_hr',\n",
       " '00000/00801_hr',\n",
       " '00000/00802_hr',\n",
       " '00000/00803_hr',\n",
       " '00000/00804_hr',\n",
       " '00000/00805_hr',\n",
       " '00000/00806_hr',\n",
       " '00000/00807_hr',\n",
       " '00000/00808_hr',\n",
       " '00000/00809_hr',\n",
       " '00000/00810_hr',\n",
       " '00000/00811_hr',\n",
       " '00000/00812_hr',\n",
       " '00000/00813_hr',\n",
       " '00000/00814_hr',\n",
       " '00000/00815_hr',\n",
       " '00000/00816_hr',\n",
       " '00000/00817_hr',\n",
       " '00000/00818_hr',\n",
       " '00000/00819_hr',\n",
       " '00000/00820_hr',\n",
       " '00000/00821_hr',\n",
       " '00000/00822_hr',\n",
       " '00000/00823_hr',\n",
       " '00000/00824_hr',\n",
       " '00000/00825_hr',\n",
       " '00000/00826_hr',\n",
       " '00000/00827_hr',\n",
       " '00000/00828_hr',\n",
       " '00000/00829_hr',\n",
       " '00000/00830_hr',\n",
       " '00000/00831_hr',\n",
       " '00000/00832_hr',\n",
       " '00000/00833_hr',\n",
       " '00000/00834_hr',\n",
       " '00000/00835_hr',\n",
       " '00000/00836_hr',\n",
       " '00000/00837_hr',\n",
       " '00000/00838_hr',\n",
       " '00000/00839_hr',\n",
       " '00000/00840_hr',\n",
       " '00000/00841_hr',\n",
       " '00000/00842_hr',\n",
       " '00000/00843_hr',\n",
       " '00000/00844_hr',\n",
       " '00000/00845_hr',\n",
       " '00000/00846_hr',\n",
       " '00000/00847_hr',\n",
       " '00000/00848_hr',\n",
       " '00000/00849_hr',\n",
       " '00000/00850_hr',\n",
       " '00000/00851_hr',\n",
       " '00000/00852_hr',\n",
       " '00000/00853_hr',\n",
       " '00000/00854_hr',\n",
       " '00000/00855_hr',\n",
       " '00000/00856_hr',\n",
       " '00000/00857_hr',\n",
       " '00000/00858_hr',\n",
       " '00000/00859_hr',\n",
       " '00000/00860_hr',\n",
       " '00000/00861_hr',\n",
       " '00000/00862_hr',\n",
       " '00000/00863_hr',\n",
       " '00000/00864_hr',\n",
       " '00000/00865_hr',\n",
       " '00000/00866_hr',\n",
       " '00000/00867_hr',\n",
       " '00000/00868_hr',\n",
       " '00000/00869_hr',\n",
       " '00000/00870_hr',\n",
       " '00000/00871_hr',\n",
       " '00000/00872_hr',\n",
       " '00000/00873_hr',\n",
       " '00000/00874_hr',\n",
       " '00000/00875_hr',\n",
       " '00000/00876_hr',\n",
       " '00000/00877_hr',\n",
       " '00000/00878_hr',\n",
       " '00000/00879_hr',\n",
       " '00000/00880_hr',\n",
       " '00000/00881_hr',\n",
       " '00000/00882_hr',\n",
       " '00000/00883_hr',\n",
       " '00000/00884_hr',\n",
       " '00000/00885_hr',\n",
       " '00000/00886_hr',\n",
       " '00000/00887_hr',\n",
       " '00000/00888_hr',\n",
       " '00000/00889_hr',\n",
       " '00000/00890_hr',\n",
       " '00000/00891_hr',\n",
       " '00000/00892_hr',\n",
       " '00000/00893_hr',\n",
       " '00000/00894_hr',\n",
       " '00000/00895_hr',\n",
       " '00000/00896_hr',\n",
       " '00000/00897_hr',\n",
       " '00000/00898_hr',\n",
       " '00000/00899_hr',\n",
       " '00000/00900_hr',\n",
       " '00000/00901_hr',\n",
       " '00000/00902_hr',\n",
       " '00000/00903_hr',\n",
       " '00000/00904_hr',\n",
       " '00000/00905_hr',\n",
       " '00000/00906_hr',\n",
       " '00000/00907_hr',\n",
       " '00000/00908_hr',\n",
       " '00000/00909_hr',\n",
       " '00000/00910_hr',\n",
       " '00000/00911_hr',\n",
       " '00000/00912_hr',\n",
       " '00000/00913_hr',\n",
       " '00000/00914_hr',\n",
       " '00000/00915_hr',\n",
       " '00000/00916_hr',\n",
       " '00000/00917_hr',\n",
       " '00000/00918_hr',\n",
       " '00000/00919_hr',\n",
       " '00000/00920_hr',\n",
       " '00000/00921_hr',\n",
       " '00000/00922_hr',\n",
       " '00000/00923_hr',\n",
       " '00000/00924_hr',\n",
       " '00000/00925_hr',\n",
       " '00000/00926_hr',\n",
       " '00000/00927_hr',\n",
       " '00000/00928_hr',\n",
       " '00000/00929_hr',\n",
       " '00000/00930_hr',\n",
       " '00000/00931_hr',\n",
       " '00000/00932_hr',\n",
       " '00000/00933_hr',\n",
       " '00000/00934_hr',\n",
       " '00000/00935_hr',\n",
       " '00000/00936_hr',\n",
       " '00000/00937_hr',\n",
       " '00000/00938_hr',\n",
       " '00000/00939_hr',\n",
       " '00000/00940_hr',\n",
       " '00000/00941_hr',\n",
       " '00000/00942_hr',\n",
       " '00000/00943_hr',\n",
       " '00000/00944_hr',\n",
       " '00000/00945_hr',\n",
       " '00000/00946_hr',\n",
       " '00000/00947_hr',\n",
       " '00000/00948_hr',\n",
       " '00000/00949_hr',\n",
       " '00000/00950_hr',\n",
       " '00000/00951_hr',\n",
       " '00000/00952_hr',\n",
       " '00000/00953_hr',\n",
       " '00000/00954_hr',\n",
       " '00000/00955_hr',\n",
       " '00000/00956_hr',\n",
       " '00000/00957_hr',\n",
       " '00000/00958_hr',\n",
       " '00000/00959_hr',\n",
       " '00000/00960_hr',\n",
       " '00000/00961_hr',\n",
       " '00000/00962_hr',\n",
       " '00000/00963_hr',\n",
       " '00000/00964_hr',\n",
       " '00000/00965_hr',\n",
       " '00000/00966_hr',\n",
       " '00000/00967_hr',\n",
       " '00000/00968_hr',\n",
       " '00000/00969_hr',\n",
       " '00000/00970_hr',\n",
       " '00000/00971_hr',\n",
       " '00000/00972_hr',\n",
       " '00000/00973_hr',\n",
       " '00000/00974_hr',\n",
       " '00000/00975_hr',\n",
       " '00000/00976_hr',\n",
       " '00000/00977_hr',\n",
       " '00000/00978_hr',\n",
       " '00000/00979_hr',\n",
       " '00000/00980_hr',\n",
       " '00000/00981_hr',\n",
       " '00000/00982_hr',\n",
       " '00000/00983_hr',\n",
       " '00000/00984_hr',\n",
       " '00000/00985_hr',\n",
       " '00000/00986_hr',\n",
       " '00000/00987_hr',\n",
       " '00000/00988_hr',\n",
       " '00000/00989_hr',\n",
       " '00000/00990_hr',\n",
       " '00000/00991_hr',\n",
       " '00000/00992_hr',\n",
       " '00000/00993_hr',\n",
       " '00000/00994_hr',\n",
       " '00000/00995_hr',\n",
       " '00000/00996_hr',\n",
       " '00000/00997_hr',\n",
       " '00000/00998_hr',\n",
       " '00000/00999_hr',\n",
       " '01000/01000_hr',\n",
       " '01000/01001_hr',\n",
       " '01000/01002_hr',\n",
       " '01000/01003_hr',\n",
       " '01000/01004_hr',\n",
       " '01000/01005_hr',\n",
       " '01000/01006_hr',\n",
       " '01000/01007_hr',\n",
       " '01000/01008_hr',\n",
       " '01000/01009_hr',\n",
       " '01000/01010_hr',\n",
       " '01000/01011_hr',\n",
       " '01000/01012_hr',\n",
       " ...]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Identify the header fiels\n",
    "record_rel_paths = find_records(args.input_dir)\n",
    "record_rel_paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f7d2f1b1-edbb-41e5-a649-e567581fe63c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(21799,\n",
       " ['00000/00001_hr', '00000/00002_hr', '00000/00003_hr', '00000/00004_hr'],\n",
       " ['21000/21834_hr', '21000/21835_hr', '21000/21836_hr', '21000/21837_hr'])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(record_rel_paths), record_rel_paths[:4], record_rel_paths[-4:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b5851faa-c9ae-400b-bff6-04a61dd07294",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 21799 records.\n"
     ]
    }
   ],
   "source": [
    "print(f\"Found {len(record_rel_paths)} records.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "0d5e3c20-bf34-4f6e-b09e-7b80d868ed7a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>RELATIVE_FILE_PATH</th>\n",
       "      <th>FILE_NAME</th>\n",
       "      <th>SAMPLE_RATE</th>\n",
       "      <th>SOURCE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Empty DataFrame\n",
       "Columns: [RELATIVE_FILE_PATH, FILE_NAME, SAMPLE_RATE, SOURCE]\n",
       "Index: []"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Prepare an index dataframe\n",
    "index_df = pd.DataFrame(columns = [\"RELATIVE_FILE_PATH\", \"FILE_NAME\", \"SAMPLE_RATE\", \"SOURCE\"])\n",
    "index_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e672e6b6-eafd-4ab2-809e-3b00523dce2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def moving_window_crop(x: np.ndarray, crop_length: int, crop_stride: int) -> np.ndarray:\n",
    "    \"\"\"Crop the input sequence with a moving window.\n",
    "    \"\"\"\n",
    "    if crop_length > x.shape[1]:\n",
    "        raise ValueError(f\"crop_length must be smaller than the length of x ({x.shape[1]}).\")\n",
    "    start_idx = np.arange(0, x.shape[1] - crop_length + 1, crop_stride)\n",
    "    return [x[:, i:i + crop_length] for i in start_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "5f8a34a1-f355-49da-a8b3-c38887eeda2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████| 21799/21799 [17:20<00:00, 20.96it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved 21799 cropped signals.\n"
     ]
    }
   ],
   "source": [
    "for record_rel_path in tqdm(record_rel_paths):\n",
    "    record_rel_dir, record_name = os.path.split(record_rel_path)\n",
    "    save_dir = os.path.join(args.output_dir, record_rel_dir)\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    source_name = record_rel_dir.split(\"/\")[0]\n",
    "    signal, record_info = wfdb.rdsamp(os.path.join(args.input_dir, record_rel_path))\n",
    "    lead_idx = np.array([record_info[\"sig_name\"].index(lead_name) for lead_name in _LEAD_NAMES])\n",
    "    signal = signal[:, lead_idx]\n",
    "    fs = record_info[\"fs\"]\n",
    "    signal_length = record_info[\"sig_len\"]\n",
    "    if signal_length < 10 * fs:  # Exclude the ECGs with lengths of less than 10 seconds\n",
    "        continue\n",
    "    cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "    for idx, cropped_signal in enumerate(cropped_signals):\n",
    "        if cropped_signal.shape[1] != 10 * fs or np.isnan(cropped_signal).any():\n",
    "            continue\n",
    "        pd.to_pickle(cropped_signal.astype(np.float32),\n",
    "                     os.path.join(save_dir, f\"{record_name}_{idx}.pkl\"))\n",
    "        index_df.loc[num_saved] = [f\"{record_rel_path}_{idx}.pkl\",\n",
    "                                   f\"{record_name}_{idx}.pkl\",\n",
    "                                   fs,\n",
    "                                   source_name]\n",
    "        num_saved += 1\n",
    "\n",
    "print(f\"Saved {num_saved} cropped signals.\")\n",
    "os.makedirs(os.path.dirname(args.index_path), exist_ok=True)\n",
    "index_df.to_csv(args.index_path, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "bf8b378c-f587-4d73-b10b-8c9e747c68da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>RELATIVE_FILE_PATH</th>\n",
       "      <th>FILE_NAME</th>\n",
       "      <th>SAMPLE_RATE</th>\n",
       "      <th>SOURCE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>00000/00001_hr_0.pkl</td>\n",
       "      <td>00001_hr_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>00000/00002_hr_0.pkl</td>\n",
       "      <td>00002_hr_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>00000/00003_hr_0.pkl</td>\n",
       "      <td>00003_hr_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>00000/00004_hr_0.pkl</td>\n",
       "      <td>00004_hr_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>00000/00005_hr_0.pkl</td>\n",
       "      <td>00005_hr_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21794</th>\n",
       "      <td>21000/21833_hr_0.pkl</td>\n",
       "      <td>21833_hr_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>21000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21795</th>\n",
       "      <td>21000/21834_hr_0.pkl</td>\n",
       "      <td>21834_hr_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>21000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21796</th>\n",
       "      <td>21000/21835_hr_0.pkl</td>\n",
       "      <td>21835_hr_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>21000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21797</th>\n",
       "      <td>21000/21836_hr_0.pkl</td>\n",
       "      <td>21836_hr_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>21000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21798</th>\n",
       "      <td>21000/21837_hr_0.pkl</td>\n",
       "      <td>21837_hr_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>21000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>21799 rows × 4 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         RELATIVE_FILE_PATH       FILE_NAME  SAMPLE_RATE SOURCE\n",
       "0      00000/00001_hr_0.pkl  00001_hr_0.pkl          500  00000\n",
       "1      00000/00002_hr_0.pkl  00002_hr_0.pkl          500  00000\n",
       "2      00000/00003_hr_0.pkl  00003_hr_0.pkl          500  00000\n",
       "3      00000/00004_hr_0.pkl  00004_hr_0.pkl          500  00000\n",
       "4      00000/00005_hr_0.pkl  00005_hr_0.pkl          500  00000\n",
       "...                     ...             ...          ...    ...\n",
       "21794  21000/21833_hr_0.pkl  21833_hr_0.pkl          500  21000\n",
       "21795  21000/21834_hr_0.pkl  21834_hr_0.pkl          500  21000\n",
       "21796  21000/21835_hr_0.pkl  21835_hr_0.pkl          500  21000\n",
       "21797  21000/21836_hr_0.pkl  21836_hr_0.pkl          500  21000\n",
       "21798  21000/21837_hr_0.pkl  21837_hr_0.pkl          500  21000\n",
       "\n",
       "[21799 rows x 4 columns]"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "628c586f-b93e-40de-a931-61003405b633",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32d2a9b9-27e5-4454-b245-b5d9fdd6c2a2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a9fd448-58b2-48c5-86e0-c2d4f632d2db",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2b6aab8-9d33-4d4c-85dc-4123ccf7ff7a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e0d95ab-d658-47d6-a21b-8693bccfa614",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f44de267-612f-49f0-a0c4-014f0a3c9066",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "145e3b32-4569-4098-b8f9-05967f08d266",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bed473c-f05b-40af-a2d5-320b009589f6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a7fe86-4137-437a-a1f5-4c6d781daef6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0751a74c-96c9-4f1c-b3d5-000636565096",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adf5d001-0cc0-4772-b232-7f84b35995ec",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d00647b-9456-4fe1-b189-a10633064f56",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e2f93b79-336d-43cb-9b07-9075e8c9d0c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                        | 0/21799 [00:00<?, ?it/s]\n"
     ]
    }
   ],
   "source": [
    "# Save all the cropped signals\n",
    "num_saved = 0\n",
    "for record_rel_path in tqdm(record_rel_paths):\n",
    "    record_rel_dir, record_name = os.path.split(record_rel_path)\n",
    "    save_dir = os.path.join(args.output_dir, record_rel_dir)\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    source_name = record_rel_dir.split(\"/\")[0]\n",
    "    signal, record_info = wfdb.rdsamp(os.path.join(args.input_dir, record_rel_path))\n",
    "    lead_idx = np.array([record_info[\"sig_name\"].index(lead_name) for lead_name in _LEAD_NAMES])\n",
    "    signal = signal[:, lead_idx]\n",
    "    fs = record_info[\"fs\"]\n",
    "    signal_length = record_info[\"sig_len\"]\n",
    "    # if signal_length < 10 * fs:  # Exclude the ECGs with lengths of less than 10 seconds\n",
    "    #     continue\n",
    "    # cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "    # for idx, cropped_signal in enumerate(cropped_signals):\n",
    "    #     if cropped_signal.shape[1] != 10 * fs or np.isnan(cropped_signal).any():\n",
    "    #         continue\n",
    "    #     pd.to_pickle(cropped_signal.astype(np.float32),\n",
    "    #                  os.path.join(save_dir, f\"{record_name}_{idx}.pkl\"))\n",
    "    #     index_df.loc[num_saved] = [f\"{record_rel_path}_{idx}.pkl\",\n",
    "    #                                f\"{record_name}_{idx}.pkl\",\n",
    "    #                                fs,\n",
    "    #                                source_name]\n",
    "    #     num_saved += 1\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "acc7bac7-3a8b-485f-a8ae-66f7cecb3ce2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('00000', '00001_hr', './ptbxl/ecgs500/00000', '00000')"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "record_rel_dir, record_name, save_dir, source_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "a1ee7f9b-5f69-493a-b3ed-2d11ec010d35",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-0.115, -0.05 ,  0.065, ..., -0.035, -0.035, -0.075],\n",
       "       [-0.115, -0.05 ,  0.065, ..., -0.035, -0.035, -0.075],\n",
       "       [-0.115, -0.05 ,  0.065, ..., -0.035, -0.035, -0.075],\n",
       "       ...,\n",
       "       [ 0.21 ,  0.205, -0.005, ...,  0.185,  0.17 ,  0.18 ],\n",
       "       [ 0.21 ,  0.205, -0.005, ...,  0.185,  0.17 ,  0.18 ],\n",
       "       [ 0.21 ,  0.205, -0.005, ...,  0.185,  0.17 ,  0.18 ]])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "67d5325a-e0d3-43d2-a882-9d00cbb02b47",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'fs': 500,\n",
       " 'sig_len': 5000,\n",
       " 'n_sig': 12,\n",
       " 'base_date': None,\n",
       " 'base_time': None,\n",
       " 'units': ['mV',\n",
       "  'mV',\n",
       "  'mV',\n",
       "  'mV',\n",
       "  'mV',\n",
       "  'mV',\n",
       "  'mV',\n",
       "  'mV',\n",
       "  'mV',\n",
       "  'mV',\n",
       "  'mV',\n",
       "  'mV'],\n",
       " 'sig_name': ['I',\n",
       "  'II',\n",
       "  'III',\n",
       "  'AVR',\n",
       "  'AVL',\n",
       "  'AVF',\n",
       "  'V1',\n",
       "  'V2',\n",
       "  'V3',\n",
       "  'V4',\n",
       "  'V5',\n",
       "  'V6'],\n",
       " 'comments': []}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "record_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aff5b601-6f4d-4004-8b15-162aa5675b4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "if signal_length < 10 * fs:  # Exclude the ECGs with lengths of less than 10 seconds\n",
    "    continue\n",
    "cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "for idx, cropped_signal in enumerate(cropped_signals):\n",
    "    if cropped_signal.shape[1] != 10 * fs or np.isnan(cropped_signal).any():\n",
    "        continue\n",
    "    pd.to_pickle(cropped_signal.astype(np.float32),\n",
    "                 os.path.join(save_dir, f\"{record_name}_{idx}.pkl\"))\n",
    "    index_df.loc[num_saved] = [f\"{record_rel_path}_{idx}.pkl\",\n",
    "                               f\"{record_name}_{idx}.pkl\",\n",
    "                               fs,\n",
    "                               source_name]\n",
    "    num_saved += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "d1521670-34dd-4150-9d70-2d43c9369593",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n"
     ]
    }
   ],
   "source": [
    "if signal_length < 10 * fs:\n",
    "    print(1)\n",
    "else: print(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "fd94d4d2-a3ef-4966-9c63-0ee79b8ca72b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5000"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "signal_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "5580e89a-0341-424a-95cc-639a29be1a6c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5000"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fs * 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "4cacee1e-f867-487e-b215-be5f364aa365",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[-0.115, -0.115, -0.115, ...,  0.21 ,  0.21 ,  0.21 ],\n",
       "        [-0.05 , -0.05 , -0.05 , ...,  0.205,  0.205,  0.205],\n",
       "        [ 0.065,  0.065,  0.065, ..., -0.005, -0.005, -0.005],\n",
       "        ...,\n",
       "        [-0.035, -0.035, -0.035, ...,  0.185,  0.185,  0.185],\n",
       "        [-0.035, -0.035, -0.035, ...,  0.17 ,  0.17 ,  0.17 ],\n",
       "        [-0.075, -0.075, -0.075, ...,  0.18 ,  0.18 ,  0.18 ]])]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "cropped_signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c86396f1-2cd3-4056-b0a0-9a1bd81042bb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5000, 12)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(signal).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "123cec3f-1d7f-4073-85e5-d03c8dfddf8c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 12, 5000)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(cropped_signals).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "7bc5c577-af73-435f-b278-b79b700baaed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-0.115 -0.115 -0.115 ...  0.21   0.21   0.21 ]\n",
      " [-0.05  -0.05  -0.05  ...  0.205  0.205  0.205]\n",
      " [ 0.065  0.065  0.065 ... -0.005 -0.005 -0.005]\n",
      " ...\n",
      " [-0.035 -0.035 -0.035 ...  0.185  0.185  0.185]\n",
      " [-0.035 -0.035 -0.035 ...  0.17   0.17   0.17 ]\n",
      " [-0.075 -0.075 -0.075 ...  0.18   0.18   0.18 ]]\n"
     ]
    }
   ],
   "source": [
    "for idx, cropped_signal in enumerate(cropped_signals):\n",
    "    print(cropped_signal)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "c15bd445-ac06-4e61-a41b-4ffc91df9395",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['00000/00001_hr_0.pkl', '00001_hr_0.pkl', 500, '00000']"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[f\"{record_rel_path}_{idx}.pkl\", f\"{record_name}_{idx}.pkl\", fs, source_name]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "81139e44-8652-4d9a-83e9-0dabab4f0c6b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>patient_id</th>\n",
       "      <th>age</th>\n",
       "      <th>sex</th>\n",
       "      <th>height</th>\n",
       "      <th>weight</th>\n",
       "      <th>nurse</th>\n",
       "      <th>site</th>\n",
       "      <th>device</th>\n",
       "      <th>recording_date</th>\n",
       "      <th>report</th>\n",
       "      <th>...</th>\n",
       "      <th>validated_by_human</th>\n",
       "      <th>baseline_drift</th>\n",
       "      <th>static_noise</th>\n",
       "      <th>burst_noise</th>\n",
       "      <th>electrodes_problems</th>\n",
       "      <th>extra_beats</th>\n",
       "      <th>pacemaker</th>\n",
       "      <th>strat_fold</th>\n",
       "      <th>filename_lr</th>\n",
       "      <th>filename_hr</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ecg_id</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>15709.0</td>\n",
       "      <td>56.0</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>63.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>CS-12   E</td>\n",
       "      <td>1984-11-09 09:17:34</td>\n",
       "      <td>sinusrhythmus periphere niederspannung</td>\n",
       "      <td>...</td>\n",
       "      <td>True</td>\n",
       "      <td>NaN</td>\n",
       "      <td>, I-V1,</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>3</td>\n",
       "      <td>records100/00000/00001_lr</td>\n",
       "      <td>records500/00000/00001_hr</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>13243.0</td>\n",
       "      <td>19.0</td>\n",
       "      <td>0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>70.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>CS-12   E</td>\n",
       "      <td>1984-11-14 12:55:37</td>\n",
       "      <td>sinusbradykardie sonst normales ekg</td>\n",
       "      <td>...</td>\n",
       "      <td>True</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>2</td>\n",
       "      <td>records100/00000/00002_lr</td>\n",
       "      <td>records500/00000/00002_hr</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>20372.0</td>\n",
       "      <td>37.0</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>69.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>CS-12   E</td>\n",
       "      <td>1984-11-15 12:49:10</td>\n",
       "      <td>sinusrhythmus normales ekg</td>\n",
       "      <td>...</td>\n",
       "      <td>True</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>5</td>\n",
       "      <td>records100/00000/00003_lr</td>\n",
       "      <td>records500/00000/00003_hr</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>17014.0</td>\n",
       "      <td>24.0</td>\n",
       "      <td>0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>82.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>CS-12   E</td>\n",
       "      <td>1984-11-15 13:44:57</td>\n",
       "      <td>sinusrhythmus normales ekg</td>\n",
       "      <td>...</td>\n",
       "      <td>True</td>\n",
       "      <td>, II,III,AVF</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>3</td>\n",
       "      <td>records100/00000/00004_lr</td>\n",
       "      <td>records500/00000/00004_hr</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>17448.0</td>\n",
       "      <td>19.0</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>70.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>CS-12   E</td>\n",
       "      <td>1984-11-17 10:43:15</td>\n",
       "      <td>sinusrhythmus normales ekg</td>\n",
       "      <td>...</td>\n",
       "      <td>True</td>\n",
       "      <td>, III,AVR,AVF</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>4</td>\n",
       "      <td>records100/00000/00005_lr</td>\n",
       "      <td>records500/00000/00005_hr</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21833</th>\n",
       "      <td>17180.0</td>\n",
       "      <td>67.0</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>AT-60    3</td>\n",
       "      <td>2001-05-31 09:14:35</td>\n",
       "      <td>ventrikulÄre extrasystole(n) sinustachykardie ...</td>\n",
       "      <td>...</td>\n",
       "      <td>True</td>\n",
       "      <td>NaN</td>\n",
       "      <td>, alles,</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1ES</td>\n",
       "      <td>NaN</td>\n",
       "      <td>7</td>\n",
       "      <td>records100/21000/21833_lr</td>\n",
       "      <td>records500/21000/21833_hr</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21834</th>\n",
       "      <td>20703.0</td>\n",
       "      <td>300.0</td>\n",
       "      <td>0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>AT-60    3</td>\n",
       "      <td>2001-06-05 11:33:39</td>\n",
       "      <td>sinusrhythmus lagetyp normal qrs(t) abnorm    ...</td>\n",
       "      <td>...</td>\n",
       "      <td>True</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>4</td>\n",
       "      <td>records100/21000/21834_lr</td>\n",
       "      <td>records500/21000/21834_hr</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21835</th>\n",
       "      <td>19311.0</td>\n",
       "      <td>59.0</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>AT-60    3</td>\n",
       "      <td>2001-06-08 10:30:27</td>\n",
       "      <td>sinusrhythmus lagetyp normal t abnorm in anter...</td>\n",
       "      <td>...</td>\n",
       "      <td>True</td>\n",
       "      <td>NaN</td>\n",
       "      <td>, I-AVR,</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>2</td>\n",
       "      <td>records100/21000/21835_lr</td>\n",
       "      <td>records500/21000/21835_hr</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21836</th>\n",
       "      <td>8873.0</td>\n",
       "      <td>64.0</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>AT-60    3</td>\n",
       "      <td>2001-06-09 18:21:49</td>\n",
       "      <td>supraventrikulÄre extrasystole(n) sinusrhythmu...</td>\n",
       "      <td>...</td>\n",
       "      <td>True</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>SVES</td>\n",
       "      <td>NaN</td>\n",
       "      <td>8</td>\n",
       "      <td>records100/21000/21836_lr</td>\n",
       "      <td>records500/21000/21836_hr</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21837</th>\n",
       "      <td>11744.0</td>\n",
       "      <td>68.0</td>\n",
       "      <td>0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>AT-60    3</td>\n",
       "      <td>2001-06-11 16:43:01</td>\n",
       "      <td>sinusrhythmus p-sinistrocardiale lagetyp norma...</td>\n",
       "      <td>...</td>\n",
       "      <td>True</td>\n",
       "      <td>NaN</td>\n",
       "      <td>, I-AVL,</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>9</td>\n",
       "      <td>records100/21000/21837_lr</td>\n",
       "      <td>records500/21000/21837_hr</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>21799 rows × 27 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "        patient_id    age  sex  height  weight  nurse  site      device  \\\n",
       "ecg_id                                                                    \n",
       "1          15709.0   56.0    1     NaN    63.0    2.0   0.0   CS-12   E   \n",
       "2          13243.0   19.0    0     NaN    70.0    2.0   0.0   CS-12   E   \n",
       "3          20372.0   37.0    1     NaN    69.0    2.0   0.0   CS-12   E   \n",
       "4          17014.0   24.0    0     NaN    82.0    2.0   0.0   CS-12   E   \n",
       "5          17448.0   19.0    1     NaN    70.0    2.0   0.0   CS-12   E   \n",
       "...            ...    ...  ...     ...     ...    ...   ...         ...   \n",
       "21833      17180.0   67.0    1     NaN     NaN    1.0   2.0  AT-60    3   \n",
       "21834      20703.0  300.0    0     NaN     NaN    1.0   2.0  AT-60    3   \n",
       "21835      19311.0   59.0    1     NaN     NaN    1.0   2.0  AT-60    3   \n",
       "21836       8873.0   64.0    1     NaN     NaN    1.0   2.0  AT-60    3   \n",
       "21837      11744.0   68.0    0     NaN     NaN    1.0   2.0  AT-60    3   \n",
       "\n",
       "             recording_date  \\\n",
       "ecg_id                        \n",
       "1       1984-11-09 09:17:34   \n",
       "2       1984-11-14 12:55:37   \n",
       "3       1984-11-15 12:49:10   \n",
       "4       1984-11-15 13:44:57   \n",
       "5       1984-11-17 10:43:15   \n",
       "...                     ...   \n",
       "21833   2001-05-31 09:14:35   \n",
       "21834   2001-06-05 11:33:39   \n",
       "21835   2001-06-08 10:30:27   \n",
       "21836   2001-06-09 18:21:49   \n",
       "21837   2001-06-11 16:43:01   \n",
       "\n",
       "                                                   report  ...  \\\n",
       "ecg_id                                                     ...   \n",
       "1                  sinusrhythmus periphere niederspannung  ...   \n",
       "2                     sinusbradykardie sonst normales ekg  ...   \n",
       "3                              sinusrhythmus normales ekg  ...   \n",
       "4                              sinusrhythmus normales ekg  ...   \n",
       "5                              sinusrhythmus normales ekg  ...   \n",
       "...                                                   ...  ...   \n",
       "21833   ventrikulÄre extrasystole(n) sinustachykardie ...  ...   \n",
       "21834   sinusrhythmus lagetyp normal qrs(t) abnorm    ...  ...   \n",
       "21835   sinusrhythmus lagetyp normal t abnorm in anter...  ...   \n",
       "21836   supraventrikulÄre extrasystole(n) sinusrhythmu...  ...   \n",
       "21837   sinusrhythmus p-sinistrocardiale lagetyp norma...  ...   \n",
       "\n",
       "       validated_by_human  baseline_drift static_noise burst_noise  \\\n",
       "ecg_id                                                               \n",
       "1                    True             NaN    , I-V1,           NaN   \n",
       "2                    True             NaN          NaN         NaN   \n",
       "3                    True             NaN          NaN         NaN   \n",
       "4                    True    , II,III,AVF          NaN         NaN   \n",
       "5                    True   , III,AVR,AVF          NaN         NaN   \n",
       "...                   ...             ...          ...         ...   \n",
       "21833                True             NaN   , alles,           NaN   \n",
       "21834                True             NaN          NaN         NaN   \n",
       "21835                True             NaN   , I-AVR,           NaN   \n",
       "21836                True             NaN          NaN         NaN   \n",
       "21837                True             NaN   , I-AVL,           NaN   \n",
       "\n",
       "        electrodes_problems  extra_beats  pacemaker  strat_fold  \\\n",
       "ecg_id                                                            \n",
       "1                       NaN          NaN        NaN           3   \n",
       "2                       NaN          NaN        NaN           2   \n",
       "3                       NaN          NaN        NaN           5   \n",
       "4                       NaN          NaN        NaN           3   \n",
       "5                       NaN          NaN        NaN           4   \n",
       "...                     ...          ...        ...         ...   \n",
       "21833                   NaN          1ES        NaN           7   \n",
       "21834                   NaN          NaN        NaN           4   \n",
       "21835                   NaN          NaN        NaN           2   \n",
       "21836                   NaN         SVES        NaN           8   \n",
       "21837                   NaN          NaN        NaN           9   \n",
       "\n",
       "                      filename_lr                filename_hr  \n",
       "ecg_id                                                        \n",
       "1       records100/00000/00001_lr  records500/00000/00001_hr  \n",
       "2       records100/00000/00002_lr  records500/00000/00002_hr  \n",
       "3       records100/00000/00003_lr  records500/00000/00003_hr  \n",
       "4       records100/00000/00004_lr  records500/00000/00004_hr  \n",
       "5       records100/00000/00005_lr  records500/00000/00005_hr  \n",
       "...                           ...                        ...  \n",
       "21833   records100/21000/21833_lr  records500/21000/21833_hr  \n",
       "21834   records100/21000/21834_lr  records500/21000/21834_hr  \n",
       "21835   records100/21000/21835_lr  records500/21000/21835_hr  \n",
       "21836   records100/21000/21836_lr  records500/21000/21836_hr  \n",
       "21837   records100/21000/21837_lr  records500/21000/21837_hr  \n",
       "\n",
       "[21799 rows x 27 columns]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y = pd.read_csv('/tf/physionet.org/files/ptb-xl/1.0.3/ptbxl_database.csv', index_col='ecg_id')\n",
    "Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "efc3e165-3730-4b01-b336-8be6a3246514",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ecg_id\n",
       "1                 {'NORM': 100.0, 'LVOLT': 0.0, 'SR': 0.0}\n",
       "2                             {'NORM': 80.0, 'SBRAD': 0.0}\n",
       "3                               {'NORM': 100.0, 'SR': 0.0}\n",
       "4                               {'NORM': 100.0, 'SR': 0.0}\n",
       "5                               {'NORM': 100.0, 'SR': 0.0}\n",
       "                               ...                        \n",
       "21833    {'NDT': 100.0, 'PVC': 100.0, 'VCLVH': 0.0, 'ST...\n",
       "21834             {'NORM': 100.0, 'ABQRS': 0.0, 'SR': 0.0}\n",
       "21835                           {'ISCAS': 50.0, 'SR': 0.0}\n",
       "21836                           {'NORM': 100.0, 'SR': 0.0}\n",
       "21837                           {'NORM': 100.0, 'SR': 0.0}\n",
       "Name: scp_codes, Length: 21799, dtype: object"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y.scp_codes # = Y.scp_codes.apply(lambda x: ast.literal_eval(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "c6c68783-152a-4e85-8cc6-6288905f0b29",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ecg_id\n",
       "1                 {'NORM': 100.0, 'LVOLT': 0.0, 'SR': 0.0}\n",
       "2                             {'NORM': 80.0, 'SBRAD': 0.0}\n",
       "3                               {'NORM': 100.0, 'SR': 0.0}\n",
       "4                               {'NORM': 100.0, 'SR': 0.0}\n",
       "5                               {'NORM': 100.0, 'SR': 0.0}\n",
       "                               ...                        \n",
       "21833    {'NDT': 100.0, 'PVC': 100.0, 'VCLVH': 0.0, 'ST...\n",
       "21834             {'NORM': 100.0, 'ABQRS': 0.0, 'SR': 0.0}\n",
       "21835                           {'ISCAS': 50.0, 'SR': 0.0}\n",
       "21836                           {'NORM': 100.0, 'SR': 0.0}\n",
       "21837                           {'NORM': 100.0, 'SR': 0.0}\n",
       "Name: scp_codes, Length: 21799, dtype: object"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ast\n",
    "Y.scp_codes = Y.scp_codes.apply(lambda x: ast.literal_eval(x))\n",
    "Y.scp_codes"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
